import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import gym
import time
import cv2
import random
import sounddevice as sd
from PIL import Image

from artear2_6.dataset import load_audio_as_spec
from artear2_6.utils import util
from artear2_6.utils.simpledaw import SimpleDAW

#####################  hyper parameters  ####################
EPISODES = 200
EP_STEPS = 200
LR_ACTOR = 0.001
LR_CRITIC = 0.002
GAMMA = 0.9
TAU = 0.01
MEMORY_CAPACITY = 10000
BATCH_SIZE = 32
RENDER = False

########################## DDPG Framework ######################
class ActorNet(nn.Module): # define the network structure for actor and critic
    def __init__(self, s_shape, a_dim):
        super(ActorNet, self).__init__()
        print(s_shape)
        self.s_dim = s_shape[0] * s_shape[1] * s_shape[2]
        self.fc1 = nn.Linear(self.s_dim, 30)
        self.fc1.weight.data.normal_(0, 0.1) # initialization of FC1
        self.out = nn.Linear(30, a_dim)
        self.out.weight.data.normal_(0, 0.1) # initilizaiton of OUT
    def forward(self, x):
        x = x.view((x.size()[0], self.s_dim))
        x = self.fc1(x)
        x = F.relu(x)
        x = self.out(x)
        x = torch.tanh(x)
        actions = x * 2 # for the game "Pendulum-v0", action range is [-2, 2]
        return actions

class CriticNet(nn.Module):
    def __init__(self, s_shape, a_dim):
        super(CriticNet, self).__init__()
        self.s_dim = s_shape[0] * s_shape[1] * s_shape[2]
        self.fcs = nn.Linear(self.s_dim, 30)
        self.fcs.weight.data.normal_(0, 0.1)
        self.fca = nn.Linear(a_dim, 30)
        self.fca.weight.data.normal_(0, 0.1)
        self.out = nn.Linear(30, 1)
        self.out.weight.data.normal_(0, 0.1)
    def forward(self, s, a):
        s = s.view((s.size()[0], self.s_dim))
        x = self.fcs(s)
        y = self.fca(a)
        actions_value = self.out(F.relu(x+y))
        return actions_value
    
class DDPG(object):
    def __init__(self, a_dim, s_shape, a_bound):
        self.a_dim, self.s_shape, self.a_bound = a_dim, s_shape, a_bound
        # self.memory = np.zeros((MEMORY_CAPACITY, s_dim * 2 + a_dim + 1), dtype=np.float32)
        self.memory = [(np.zeros(s_shape), np.zeros(a_dim), [0.], np.zeros(s_shape))] * MEMORY_CAPACITY
        self.pointer = 0 # serves as updating the memory data 
        # Create the 4 network objects
        self.actor_eval = ActorNet(s_shape, a_dim)
        self.actor_target = ActorNet(s_shape, a_dim)
        self.critic_eval = CriticNet(s_shape, a_dim)
        self.critic_target = CriticNet(s_shape, a_dim)
        # create 2 optimizers for actor and critic
        self.actor_optimizer = torch.optim.Adam(self.actor_eval.parameters(), lr=LR_ACTOR)
        self.critic_optimizer = torch.optim.Adam(self.critic_eval.parameters(), lr=LR_CRITIC)
        # Define the loss function for critic network update
        self.loss_func = nn.MSELoss()
        
    def store_transition(self, s, a, r, s_): # how to store the episodic data to buffer
        transition = [(s, a, [r], s_)]
        index = self.pointer % MEMORY_CAPACITY # replace the old data with new data 
        self.memory[index] = transition
        self.pointer += 1
    
    def choose_action(self, s):
        # print(s)
        s = torch.unsqueeze(torch.FloatTensor(s), 0)
        print(s.size())
        return self.actor_eval(s)[0].detach()
    
    def learn(self):
        # softly update the target networks
        for x in self.actor_target.state_dict().keys():
            eval('self.actor_target.' + x + '.data.mul_((1-TAU))')
            eval('self.actor_target.' + x + '.data.add_(TAU*self.actor_eval.' + x + '.data)')
        for x in self.critic_target.state_dict().keys():
            eval('self.critic_target.' + x + '.data.mul_((1-TAU))')
            eval('self.critic_target.' + x + '.data.add_(TAU*self.critic_eval.' + x + '.data)')           
        # sample from buffer a mini-batch data 从小批量数据中采样
        indices = np.random.choice(MEMORY_CAPACITY, size=BATCH_SIZE)
        batch_trans = self.memory[indices]
        # extract data from mini-batch of transitions including s, a, r, s_ 
        
        batch_s, batch_a, batch_r, batch_s_ = [], [], [], []
        for trans in batch_trans:
            batch_s.append(trans[0])
            batch_a.append(trans[1])
            batch_r.append(trans[2])
            batch_s_.append(trans[3])
        batch_s = torch.FloatTensor(batch_s)
        batch_a = torch.FloatTensor(batch_a)
        batch_r = torch.FloatTensor(batch_r)
        batch_s_ = torch.FloatTensor(batch_s_)
        
        # make action and evaluate its action values
        a = self.actor_eval(batch_s)
        q = self.critic_eval(batch_s, a)
        actor_loss = -torch.mean(q)
        # optimize the loss of actor network
        self.actor_optimizer.zero_grad()
        actor_loss.backward()
        self.actor_optimizer.step()
        
        # compute the target Q value using the information of next state
        a_target = self.actor_target(batch_s_)
        q_tmp = self.critic_target(batch_s_, a_target)
        q_target = batch_r + GAMMA * q_tmp
        # compute the current q value and the loss
        q_eval = self.critic_eval(batch_s, batch_a)
        td_error = self.loss_func(q_target, q_eval)
        # optimize the loss of critic network
        self.critic_optimizer.zero_grad()
        td_error.backward()
        self.critic_optimizer.step()
        

# 通过得到RGB每个通道的直方图来计算相似度
def classify_hist_with_split(image1, image2, size=(256, 256)):
    # 将图像resize后，分离为RGB三个通道，再计算每个通道的相似值
    image1 = cv2.resize(image1, size)
    image2 = cv2.resize(image2, size)
    sub_image1 = cv2.split(image1)
    sub_image2 = cv2.split(image2)
    sub_data = 0
    for im1, im2 in zip(sub_image1, sub_image2):
        sub_data += calculate(im1, im2)
    sub_data = sub_data / 3
    return sub_data


# 计算单通道的直方图的相似值
def calculate(image1, image2):
    hist1 = cv2.calcHist([image1], [0], None, [256], [0.0, 255.0])
    hist2 = cv2.calcHist([image2], [0], None, [256], [0.0, 255.0])
    # 计算直方图的重合度
    degree = 0
    for i in range(len(hist1)):
        if hist1[i] != hist2[i]:
            degree = degree + (1 - abs(hist1[i] - hist2[i]) / max(hist1[i], hist2[i]))
        else:
            degree = degree + 1
    degree = degree / len(hist1)
    return degree
        
        
class Simulator:
    def __init__(self):
        self.daw = SimpleDAW(plugin='C:/VST/64bit/Sylenth1.dll', sample_rate=44100, bpm=120)
        self.midi_paths = util.get_files(r'B:\muse_repo\MIDI', 'mid')
        
        self.daw.load_midi(random.choice(self.midi_paths))
        # self.daw.set_params(list(action))
        audio = self.daw.render(4.)
        print(audio.shape)
        
        self.action_space = np.zeros(244)
        self.observation_space = load_audio_as_spec('', 4., 44100, 50., raw=audio)[0]
        self.state_shape = self.observation_space.shape
        self.action_space_high = 1.
        self.action_space_low = 0.
        self.step_count = 0

    def step(self, action):
        # self.daw.load_midi(random.choice(self.midi_paths))
        self.daw.set_params(list(action))
        print('正在播放预测声音')
        audio = self.daw.render(4.)
        state = load_audio_as_spec('', 4., 44100, 50., raw=audio)[0]
        sd.play(audio, 44100, blocking=False)
        print('播放结束')
        
        print(state.shape)
        
        print(np.uint8(self.observation_space)[:10, ...])
        _obs_space = Image.fromarray(np.uint8(self.observation_space))
        _obs_space.convert('P')
        _obs_space = np.array(_obs_space)
        _state = Image.fromarray(state)
        _state.convert('P')
        _state = np.array(_state)
        
        reward = classify_hist_with_split(_obs_space, _state, size=state.shape[:2])
        print('reward:', reward)
        
        self.step_count += 1
        print('[第{}轮训练]'.format(self.step_count), 'action:', action[:10], 'state:', state[:10])
        
        done = True
        info = {}
        return state, reward, done, info

    def reset(self):
        state = np.zeros(self.state_shape)
        return state

    def render(self, mode='human'):
        pass

    def seed(self, seed=None):
        pass
        
############################### Training ######################################
# Define the env in gym
env = Simulator()
s_shape = env.observation_space.shape
a_dim = env.action_space.shape[0]
a_bound = env.action_space_high
a_low_bound = env.action_space_low

ddpg = DDPG(a_dim, s_shape, a_bound)
var = 3 # the controller of exploration which will decay during training process
t1 = time.time()
for i in range(EPISODES):
    s = env.reset()
    ep_r = 0
    for j in range(EP_STEPS):
        if RENDER: env.render()
        # add explorative noise to action
        a = ddpg.choose_action(s)
        a = np.clip(np.random.normal(a, var), a_low_bound, a_bound)
        s_, r, done, info = env.step(a)
        ddpg.store_transition(s, a, r / 10, s_) # store the transition to memory
        
        if ddpg.pointer > MEMORY_CAPACITY:
            var *= 0.9995 # decay the exploration controller factor
            ddpg.learn()
            
        s = s_
        ep_r += r
        if j == EP_STEPS - 1:
            print('Episode: ', i, ' Reward: %i' % (ep_r), 'Explore: %.2f' % var)
            if ep_r > -300 : RENDER = True
            break
print('Running time: ', time.time() - t1)
